import asyncio
import os

import openai
from langchain.schema import HumanMessage

from a0base.base_llm import pop_llm

openai.api_key = os.getenv("OPENAI_API_KEY")


class MQ:
    """最简陋的内存消息队列"""

    def __init__(self): self.q = asyncio.Queue()

    async def send(self, msg): await self.q.put(msg)

    async def recv(self): return await self.q.get()


mq = MQ()


# ---- Agent A：诗人 ----
async def poet_loop():
    while True:
        msg = await mq.recv()
        if msg == "exit": break
        resp = await pop_llm.ainvoke([HumanMessage(content=f"写一句七言诗，上一句是：{msg}")])
        print(f"[Poet ] {resp.content}")
        await mq.send(resp.content)


# ---- Agent B：评论家 ----
async def critic_loop():
    await asyncio.sleep(1)  # 让诗人先写
    seed = "春风得意马蹄疾"
    await mq.send(seed)
    while True:
        poem = await mq.recv()
        if poem == "exit": break
        review = await pop_llm.ainvoke([HumanMessage(content=f"用 1 句话点评这句诗：{poem}")])
        print(f"[Critic] {review.content}")
        await asyncio.sleep(2)
        await mq.send(review.content)  # 把点评再发回给诗人


async def main():
    await asyncio.gather(poet_loop(), critic_loop())


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        asyncio.run(mq.send("exit"))
